package union_find

type UnionFindV3 struct {
	parent []int
	sz     []int
}

func (u *UnionFindV3) Size() int {
	return len(u.parent)
}

func (u *UnionFindV3) find(p int) int {
	if p < 0 || p > len(u.parent) {
		panic("p is out of range.")
	}

	// 不断去查询自己的父亲节点, 直到到达根节点
	// 根节点的特点: parent[p] == p
	for p != u.parent[p] {
		p = u.parent[p]
	}
	return p
}

func (u *UnionFindV3) IsConnected(p int, q int) bool {
	return u.find(p) == u.find(q)
}

func (u *UnionFindV3) Union(p int, q int) {
	pRoot := u.find(p)
	qRoot := u.find(q)

	if pRoot == qRoot {
		return
	}

	// 根据两个元素所在树的元素个数不同判断合并方向
	// 将元素个数少的集合合并到元素个数多的集合上
	if u.sz[pRoot] < u.sz[qRoot] {
		u.parent[pRoot] = qRoot
		u.sz[qRoot] += u.sz[pRoot]
	} else {
		u.parent[qRoot] = pRoot
		u.sz[pRoot] += u.sz[qRoot]
	}

}

func NewUFV3(size int) *UnionFindV3 {
	sz := make([]int, size)
	parent := make([]int, size)

	for i := 0; i < len(parent); i++ {
		sz[i] = 1
		parent[i] = i
	}
	return &UnionFindV3{
		parent: parent,
		sz:     sz,
	}
}
